{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Debug Notebook\n",
    "23-11-23"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '1'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax import random\n",
    "from jax import jit, vmap\n",
    "from flash_attention_jax import flash_attention\n",
    "\n",
    "rng_key = random.PRNGKey(42)\n",
    "\n",
    "q = random.normal(rng_key, (1, 2, 131072, 512))  # (batch, heads, seq, dim)\n",
    "k = random.normal(rng_key, (1, 2, 131072, 512))\n",
    "v = random.normal(rng_key, (1, 2, 131072, 512))\n",
    "mask = random.randint(rng_key, (1, 131072,), 0, 2) # (batch, seq)\n",
    "\n",
    "%timeit jit(flash_attention)(q, k, v, mask).block_until_ready()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit flash_attention(q, k, v, mask).block_until_ready()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from flash_attention_jax import plain_attention, flash_attention, value_and_grad_difference\n",
    "\n",
    "diff, (dq_diff, dk_diff, dv_diff) = value_and_grad_difference(\n",
    "    plain_attention,\n",
    "    flash_attention,\n",
    "    seed = 42\n",
    ")\n",
    "\n",
    "print('shows differences between normal and flash attention for output, dq, dk, dv')\n",
    "print(f'o: {diff}')       # < 1e-4\n",
    "print(f'dq: {dq_diff}')   # < 1e-6\n",
    "print(f'dk: {dk_diff}')   # < 1e-6\n",
    "print(f'dv: {dv_diff}')   # < 1e-6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax import random\n",
    "from flash_attention_jax import causal_flash_attention\n",
    "\n",
    "rng_key = random.PRNGKey(42)\n",
    "\n",
    "q = random.normal(rng_key, (1, 2, 131072, 512)) \n",
    "k = random.normal(rng_key, (1, 2, 131072, 512))\n",
    "v = random.normal(rng_key, (1, 2, 131072, 512))\n",
    "\n",
    "out, _ = causal_flash_attention(q, k, v)\n",
    "\n",
    "out.shape  # (131072, 512)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "import haiku as hk\n",
    "import jax.numpy as jnp\n",
    "\n",
    "# example\n",
    "class Net(hk.Module):\n",
    "\n",
    "\tdef __init__(self,\n",
    "\t\t\t     dim_out: int,\n",
    "\t\t\t\t name=\"net\",):\n",
    "\t\tsuper().__init__(name=name)\n",
    "\n",
    "\t\tself.linear = hk.Linear(output_size=dim_out, with_bias=True, name=\"linear_with_bias\")\n",
    "\t\tself.act_fn = jax.nn.relu\n",
    "\t\n",
    "\tdef __call__(self, x):\n",
    "\n",
    "\t\tn_channel = x.shape[-1]\n",
    "\t\tvec_in = jnp.reshape(x, (-1, n_channel))\n",
    "\t\tvec_out = self.act_fn(self.linear(vec_in))\n",
    "\t\treturn vec_out\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import mindspore as ms\n",
    "import mindspore.nn as nn\n",
    "from mindspore.ops import functional as F\n",
    "\n",
    "class Net(nn.Cell):\n",
    "\n",
    "    def __init__(self,\n",
    "                 dim_in: int,\n",
    "                 dim_out: int,\n",
    "                 ):\n",
    "        super().__init__()\n",
    "\n",
    "        self.dense = nn.Dense(in_channels=dim_in, \n",
    "                               out_channels=dim_out, \n",
    "                               has_bias=True, \n",
    "                               activation=nn.ReLU())\n",
    "    \n",
    "    def construct(self, x):\n",
    "\n",
    "        n_channel = x.shape[-1]\n",
    "        vec_in = F.reshape(x, (-1, n_channel))\n",
    "        vec_out = self.dense(vec_in)\n",
    "        return vec_out"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "jax_ms",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
